#include "StdAfx.h"
#include "TcpClient.h"

#include <sstream>

#define WIN32_LEAN_AND_MEAN
#define _HAS_EXCEPTIONS 1
#include <windows.h>

#include <stdlib.h>
#include <stdio.h>
#include "conio.h"

#include <cctype>

#include <typeinfo>
using namespace std;

// Need to link with Ws2_32.lib, Mswsock.lib, and Advapi32.lib
#pragma comment (lib, "Ws2_32.lib")
#pragma comment (lib, "Mswsock.lib")
#pragma comment (lib, "AdvApi32.lib")

#define DEFAULT_BUFLEN 512

namespace libdmail
{

	TcpClient::TcpClient(void):mConnectSocket(INVALID_SOCKET), mServer(""), mPort(0)
	{
	}

	TcpClient::TcpClient(Protocol *protocol)
	{
		this->mProtocol = protocol;
	}

	TcpClient::~TcpClient(void)
	{
	}

	
	int TcpClient::connectServer(string server, string port)
	{
		this->mServer = server;
		this->mPort = port;

		WSADATA wsaData;
	
		struct addrinfo *result = NULL,
						*ptr = NULL,
						hints;
	
		int iResult;
	
		// Initialize Winsock
		iResult = WSAStartup(MAKEWORD(2,2), &wsaData);
		if (iResult != 0) {
			printf("WSAStartup failed with error: %d\n", iResult);
			return 1;
		}

		ZeroMemory( &hints, sizeof(hints) );
		hints.ai_family = AF_UNSPEC;
		hints.ai_socktype = SOCK_STREAM;
		hints.ai_protocol = IPPROTO_TCP;

		// Resolve the server address and port
		iResult = getaddrinfo(this->mServer.c_str(), this->mPort.c_str(), &hints, &result);
		if ( iResult != 0 ) {
			printf("getaddrinfo failed with error: %d\n", iResult);
			WSACleanup();
			return 1;
		}
		
		// Attempt to connect to an address until one succeeds
		for(ptr=result; ptr != NULL ;ptr=ptr->ai_next) {

			// Create a SOCKET for connecting to server
			mConnectSocket = socket(ptr->ai_family, ptr->ai_socktype, 
				ptr->ai_protocol);
			if (mConnectSocket == INVALID_SOCKET) {
				printf("socket failed with error: %ld\n", WSAGetLastError());
				WSACleanup();
				return 1;
			}
	 
			// Connect to server.
			iResult = connect( mConnectSocket, ptr->ai_addr, (int)ptr->ai_addrlen);
			if (iResult == SOCKET_ERROR) {
				closesocket(mConnectSocket);
				mConnectSocket = INVALID_SOCKET;
				continue;
			}
			break;
		}

		unsigned long iNonBlock = 1;
		iResult = ioctlsocket(mConnectSocket, FIONBIO, &iNonBlock);

		if (iResult != NO_ERROR)
		{
			printf("ioctlsocket failed with error: %ld\n", iResult);
		}

		//freeaddrinfo(result);

		if (mConnectSocket == INVALID_SOCKET) {
			printf("Unable to connect to server!\n");
			WSACleanup();
			return 1;
		}

		return 0;
	}

	int TcpClient::disconnectServer()
	{
		return 0;
	}

	int TcpClient::receiveData()
	{
		int iResult;
		fd_set readfds;
		struct timeval tv;

		tv.tv_sec = 0;
		tv.tv_usec = 0;
		    
		char recvbuf[DEFAULT_BUFLEN+1];
		int recvbuflen = DEFAULT_BUFLEN;
	    
		char recvCache[DEFAULT_BUFLEN];
		memset(recvCache, 0, DEFAULT_BUFLEN);

		stringstream strStream(stringstream::in | stringstream::out);
 
		do {
			FD_ZERO(&readfds);
		
			FD_SET(mConnectSocket, &readfds);
			
			// don't care about writefds and exceptfds:
			int i = select(mConnectSocket+1, &readfds, NULL, NULL, &tv);

			if (0 < i)
			{
				if (FD_ISSET(mConnectSocket, &readfds))
				{
					iResult = recv(mConnectSocket, recvbuf, recvbuflen, 0);

					if ( iResult > 0 ) 
					{
						recvbuf[iResult] = '\0';

						//strcat(recvCache,recvbuf);
 
						strStream.write(recvbuf, iResult);
						
						printf("S:\r\n%s", recvbuf);	

						int iStrPos = strStream.tellp();
 
						int iRet = this->mProtocol->onDataReceived(&strStream, iStrPos);

						if (-1 == iRet)
						{
							printf("Quit\n");
							break;
						}
					}
					else if ( iResult == 0 ) {
						printf("Connection closed\n");
						break;
					}
					else {
						printf("recv failed with error: %d\n", WSAGetLastError());
						break;
					}
				}
			}

			if (SOCKET_ERROR == i)
				printf("Socket failed with error in select: %d\n", WSAGetLastError());
		 
		} while(1);
		 

		// shutdown the connection since no more data will be sent
		iResult = shutdown(mConnectSocket, SD_SEND);
		if (iResult == SOCKET_ERROR) {
			printf("shutdown failed with error: %d\n", WSAGetLastError());
			closesocket(mConnectSocket);
			WSACleanup();
			return 1;
		}
	 
		// cleanup
		closesocket(mConnectSocket);
		WSACleanup();

	}

	int TcpClient::sendData(string data)
	{
		if (INVALID_SOCKET  == mConnectSocket)
		{
			this->connectServer(this->mServer, this->mPort);
		}

		int iResult;
		fd_set writefds;
		struct timeval tv;

		tv.tv_sec = 0;
		tv.tv_usec = 0;
	
		do {
			FD_ZERO(&writefds);
			FD_SET(mConnectSocket, &writefds);
	 
			// don't care about readfds and exceptfds:
			int i = select(mConnectSocket+1, NULL, &writefds, NULL, &tv);

			if (0 < i)
			{
				if (FD_ISSET(mConnectSocket, &writefds))
				{ 
					// Send an initial buffer
					iResult = send( mConnectSocket, data.c_str(), (int)data.length(), 0 );
					
					printf("C:%s", data.c_str());

					if (iResult == SOCKET_ERROR) {
						printf("send failed with error: %d\n", WSAGetLastError());
						closesocket(mConnectSocket);
						WSACleanup();
						return 1;
					}

					return 0;
				}
			}

			if (SOCKET_ERROR == i)
				printf("Socket failed with error in select: %d\n", WSAGetLastError());
		 
		} while(1);
		 

		// shutdown the connection since no more data will be sent
		iResult = shutdown(mConnectSocket, SD_SEND);
		if (iResult == SOCKET_ERROR) {
			printf("shutdown failed with error: %d\n", WSAGetLastError());
			closesocket(mConnectSocket);
			WSACleanup();
			return 1;
		}
	 
		// cleanup
		closesocket(mConnectSocket);
		WSACleanup();

	}

}